import jieba  # 导入 jieba 中文分词库
# 导入 defaultdict 和 Counter，用于构建频率统计表
from collections import defaultdict, Counter

# 示例语料库（注意：这些是连续的中文句子，没有空格）
corpus = [
    "我喜欢吃苹果",
    "我喜欢吃香蕉",
    "她喜欢吃葡萄",
    "他不喜欢吃香蕉",
    "他喜欢吃苹果",
    "她喜欢吃草莓"
]

# 1. 构建 Bigram 统计表，格式为 bigrams[前一个词][下一个词] = 出现次数
bigrams = defaultdict(Counter)  # 自动初始化嵌套结构，避免手动判断键是否存在

# 遍历每一句话
for sentence in corpus:
    words = list(jieba.cut(sentence))  # 使用 jieba 对句子进行分词，返回词列表
    # 构建二元组（Bigram）：每两个连续词之间的搭配
    for i in range(len(words) - 1):
        w1, w2 = words[i], words[i + 1]  # 当前词和下一个词
        bigrams[w1][w2] += 1  # 出现一次就 +1

# 2. 将频率转换成概率，构建 Bigram 概率模型
bigram_probs = {}  # 最终保存的是 P(w2 | w1) 的概率表

# 遍历所有前词（w1）
for w1 in bigrams:
    total = sum(bigrams[w1].values())  # w1 后面出现所有词的总次数
    # 计算每一个 w2 的条件概率：P(w2 | w1) = 频率 / 总数
    bigram_probs[w1] = {
        w2: count / total for w2, count in bigrams[w1].items()
    }

# 3. 定义一个函数：输入一个词，预测它后面最可能出现的词（按概率从大到小排序）


def predict_next_word(word):
    if word not in bigram_probs:
        return "未知（没有数据）"  # 如果这个词从未出现在训练语料中
    # 按照概率从高到低排序返回
    sorted_probs = sorted(
        bigram_probs[word].items(), key=lambda x: x[1], reverse=True
    )
    return sorted_probs


# 4. 打印所有 Bigram 的条件概率
print("=== Bigram 概率 ===")
for w1 in bigram_probs:
    for w2 in bigram_probs[w1]:
        prob = bigram_probs[w1][w2]
        print(f"P({w2} | {w1}) = {prob:.2f}")  # 输出形如：P(吃 | 喜欢) = 1.00

# 5. 测试预测某些词后面最可能出现的词
print("\n=== 预测示例 ===")
for test_word in ["我", "吃", "喜欢", "苹果"]:
    print(f"{test_word} → {predict_next_word(test_word)}")